#include "core/main.h"

void NeighborhoodCollection::add(const Neighborhood &neigh)
{
    // cout << "Called add" << endl;
    counts[neigh]++;

    // Print everything for fun
    // for (auto it = counts.begin(); it != counts.end(); ++it) {
    //     cout << "key: " << it->first << " value: " << it->second << endl;
    // }
}

void NeighborhoodCollection::addAll(const FrameGrid &frameGrid, int size)
{
    for (int y = 0; y + size < frameGrid.height(); y++) {
        for (int x = 0; x + size < frameGrid.width(); x++) {
            add(frameGrid.getNeighborhood(x, y, size));
        }
    }
}

struct CompareSecondDescend {
    template<class p>
    bool operator()(const p &lhs, const p &rhs) {
        return lhs.second > rhs.second;
    }
};

void NeighborhoodCollection::getMostFrequent(int num, set<Neighborhood> &out) const {
    set<pair<Neighborhood, int> > withCounts;
    getMostFrequentWithCount(num, withCounts);
    for (auto it = withCounts.begin(); it != withCounts.end(); ++it) {
        out.insert(it->first);
    }
}

void NeighborhoodCollection::getMostFrequentWithCount(int num, set<pair<Neighborhood, int> > &out) const
{
    // cout << "Number of distinct neighborhoods: " << counts.size() << endl;
    vector<pair<Neighborhood, int> > vec;

    // Add all pairs to the vector
    for (auto it = counts.begin(); it != counts.end(); ++it) {
        vec.push_back(*it);
    }

    // Sort by the second element
    CompareSecondDescend comp;
    partial_sort(vec.begin(), vec.begin() + num, vec.end(), comp);
    vec.resize(num);

    // Add the results to the set
    for (auto it = vec.begin(); it != vec.end(); ++it) {
        out.insert(*it);
    }
}